from te import tik

def custom_add():
    tik_instance = tik.Tik()

    # 设定 shape 和数据类型
    shape = (1024,)
    dtype = "float16"

    # 分配张量到 GM（全局内存）与 UB（统一缓冲区）
    input_x = tik_instance.Tensor(dtype, shape, name="input_x", scope=tik.scope_gm)
    input_y = tik_instance.Tensor(dtype, shape, name="input_y", scope=tik.scope_gm)
    output = tik_instance.Tensor(dtype, shape, name="output", scope=tik.scope_gm)

    ub_x = tik_instance.Tensor(dtype, shape, name="ub_x", scope=tik.scope_ubuf)
    ub_y = tik_instance.Tensor(dtype, shape, name="ub_y", scope=tik.scope_ubuf)
    ub_res = tik_instance.Tensor(dtype, shape, name="ub_res", scope=tik.scope_ubuf)

    # 复制输入数据到 UB
    tik_instance.data_move(ub_x, input_x, 0, 1, 1024 // 16, 0, 0)
    tik_instance.data_move(ub_y, input_y, 0, 1, 1024 // 16, 0, 0)

    # 元素加法
    with tik_instance.for_range(0, 1024) as i:
        ub_res[i] = ub_x[i] + ub_y[i]

    # 将结果写回 GM
    tik_instance.data_move(output, ub_res, 0, 1, 1024 // 16, 0, 0)

    return tik_instance.BuildCCE(kernel_name="custom_add", inputs=[input_x, input_y], outputs=[output])